package com.qhhc.websocket;

import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;


public class VoiceWebSocketHandler extends AbstractWebSocketHandler {
    private static final Logger logger = LoggerFactory.getLogger(VoiceWebSocketHandler.class);
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private final Gson gson = new Gson();

//    @Autowired
//    private VoiceService voiceService;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String sessionId = UUID.randomUUID().toString();
        sessions.put(sessionId, session);
        logger.info("WebSocket连接已建立，会话ID: {}", sessionId);

        // 发送会话ID给客户端
        JsonObject response = new JsonObject();
        response.addProperty("type", "connection");
        response.addProperty("sessionId", sessionId);
        sendMessage(session, response.toString());
    }


    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        logger.error("WebSocket传输错误", exception);
        sendError(session, "传输错误: " + exception.getMessage());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String sessionId = sessions.entrySet().stream()
                .filter(entry -> entry.getValue().equals(session))
                .map(Map.Entry::getKey)
                .findFirst()
                .orElse(null);

        if (sessionId != null) {
            sessions.remove(sessionId);
            logger.info("WebSocket连接已关闭，会话ID: {}, 状态码: {}, 原因: {}",
                    sessionId, status.getCode(), status.getReason());
        }
    }

    private void sendMessage(WebSocketSession session, String message) {
        try {
            if (session.isOpen()) {
                session.sendMessage(new TextMessage(message));
            }
        } catch (Exception e) {
            logger.error("发送WebSocket消息时出错", e);
        }
    }

    private void sendError(WebSocketSession session, String errorMessage) {
        JsonObject error = new JsonObject();
        error.addProperty("type", "error");
        error.addProperty("message", errorMessage);
        sendMessage(session, error.toString());
    }

    // 添加一个新的Map来存储元数据
    private final Map<String, JsonObject> pendingMetadata = new ConcurrentHashMap<>();
    private final OkHttpClient httpClient = new OkHttpClient();

    @Value("${sensevoice.local.url:http://localhost:8000/api/v1/asr}")
    private String localSenseVoiceUrl;

    // 添加二进制消息处理方法
    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        String sessionId = getSessionId(session);
        logger.info("收到WebSocket二进制消息，大小: {} 字节, 会话ID: {}", message.getPayloadLength(), sessionId);

        // 获取之前保存的元数据
        JsonObject metadata = pendingMetadata.get(sessionId);
        if (metadata == null) {
            logger.error("收到二进制数据但没有对应的元数据，会话ID: {}", sessionId);
            sendError(session, "收到二进制数据但没有对应的元数据");
            return;
        }

        try {
            // 提取元数据信息
            boolean useLocalStorage = metadata.get("useLocalStorage").getAsBoolean();
            boolean isFinal = metadata.get("isFinal").getAsBoolean();
            String format = metadata.has("format") ? metadata.get("format").getAsString() : "audio/webm";

            logger.info("处理二进制音频数据，大小: {} 字节, 格式: {}, 使用本地存储: {}, 是否为最终数据: {}",
                    message.getPayloadLength(), format, useLocalStorage, isFinal);

            // 获取二进制数据
            ByteBuffer buffer = message.getPayload();
            byte[] audioData = new byte[buffer.remaining()];
            buffer.get(audioData);

            // 处理结果
            String result="";

            if (useLocalStorage) {
                // 直接发送给SenseVoice API
                result = processWithSenseVoice(audioData, format);
            } else {
                // 使用阿里云服务
//                result = voiceService.processAudio(
//                        new VoiceService.CustomMultipartFile(
//                                audioData,
//                                "audio." + getExtensionFromFormat(format),
//                                "audio." + getExtensionFromFormat(format),
//                                format),
//                        false, // 不使用本地存储
//                        isFinal);
            }

            // 发送识别结果
            JsonObject response = new JsonObject();
            response.addProperty("type", "recognition");
            response.addProperty("text", result);
            response.addProperty("isFinal", isFinal);
            sendMessage(session, response.toString());

            // 处理完成后移除元数据（仅对最终数据）
            if (isFinal) {
                pendingMetadata.remove(sessionId);
            }
        } catch (Exception e) {
            logger.error("处理WebSocket二进制消息时出错", e);
            sendError(session, "处理二进制消息时出错: " + e.getMessage());
            pendingMetadata.remove(sessionId);
        }
    }

    private String processWithSenseVoice(byte[] audioData, String format) {
        try {
            // 获取SenseVoice服务二进制API URL
            String senseVoiceUrl = localSenseVoiceUrl;
            if (!senseVoiceUrl.endsWith("/")) {
                senseVoiceUrl += "/";
            }
            // 使用新的直接二进制端点
            senseVoiceUrl += "binary_direct";
            // 添加查询参数
            senseVoiceUrl += "?lang=zh&format=" + format;

            logger.info("使用SenseVoice二进制API处理数据，URL: {}, 大小: {} bytes, 格式: {}",
                    senseVoiceUrl, audioData.length, format);

            // 直接发送二进制数据
            RequestBody requestBody = RequestBody.create(
                    MediaType.parse("application/octet-stream"),
                    audioData
            );

            Request request = new Request.Builder()
                    .url(senseVoiceUrl)
                    .post(requestBody)
                    .build();

            try (Response response = httpClient.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    throw new IOException("本地SenseVoice服务返回错误: " + response.code());
                }

                String responseBody = response.body().string();
                logger.info("SenseVoice二进制API响应: {}", responseBody);

                // 解析JSON响应
                JsonObject jsonResponse = gson.fromJson(responseBody, JsonObject.class);

                // 检查是否有错误
                if (jsonResponse.has("error")) {
                    logger.error("SenseVoice返回错误: {}", jsonResponse.get("error").getAsString());
                    return "未能识别语音内容";
                }

                // 提取识别结果
                JsonArray resultArray = jsonResponse.getAsJsonArray("result");
                if (resultArray != null && resultArray.size() > 0) {
                    JsonObject firstResult = resultArray.get(0).getAsJsonObject();
                    return firstResult.get("text").getAsString();
                } else {
                    return "未能识别语音内容";
                }
            }
        } catch (Exception e) {
            logger.error("使用SenseVoice二进制API处理数据失败", e);
            throw new RuntimeException("使用SenseVoice二进制API处理数据失败: " + e.getMessage());
        }
    }

    private String getExtensionFromFormat(String format) {
        if (format == null) return "webm";

        if (format.contains("webm")) return "webm";
        if (format.contains("wav")) return "wav";
        if (format.contains("mp3")) return "mp3";
        if (format.contains("ogg")) return "ogg";

        // 默认返回webm
        return "webm";
    }

    // 修改文本消息处理方法，添加元数据处理
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        try {
            String payload = message.getPayload();
            logger.info("收到WebSocket文本消息，长度: {}", payload.length());

            JsonObject request = gson.fromJson(payload, JsonObject.class);
            String type = request.get("type").getAsString();

          if ("audio_binary".equals(type) || "audio_metadata".equals(type)) {
                // 保存元数据，等待二进制数据
                String sessionId = getSessionId(session);
                pendingMetadata.put(sessionId, request);
                logger.info("收到音频元数据，等待二进制数据, 会话ID: {}, 格式: {}",
                        sessionId,
                        request.has("format") ? request.get("format").getAsString() : "未知");
            }
        } catch (Exception e) {
            logger.error("处理WebSocket文本消息时出错", e);
            sendError(session, "处理消息时出错: " + e.getMessage());
        }
    }

    // 辅助方法：获取会话ID
    private String getSessionId(WebSocketSession session) {
        return sessions.entrySet().stream()
                .filter(entry -> entry.getValue().equals(session))
                .map(Map.Entry::getKey)
                .findFirst()
                .orElse("unknown");
    }
}